import os
import gc
import re
import time
import uuid
import argparse
import datetime
import sys
import numpy as np
from tqdm import tqdm

import torch
import torch.nn.functional as F
from torch_sparse import SparseTensor
from torch_sparse import remove_diag, set_diag

from hgb.params import parse_args
from dataProcess import *

# 添加 se-hgnn 目录到 sys.path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))


from model import *
from utils import *
from sparseTools.utils import SparseAdjList


def main(args):
    # 加载数据集，返回图g、邻接矩阵adjs、初始标签init_labels、类别数量num_classes、数据加载器dl、训练验证节点trainval_nid和测试节点test_nid。
    global num_tgt_nodes
    g, adjs, init_labels, num_classes, dl, trainval_nid, test_nid = load_dataset(args)

    # 这句代码的整体作用是将邻接矩阵中每条边的权重重新设置为该边起始节点度的倒数。这样做的目的是进行归一化处理，使得每个节点连接的边权重之和为1，通常用于消息传递中的归一化步骤。
    adjs = adjsProcess(adjs)

    # =======
    # neighbor aggregation
    # =======
    # 根据数据集的不同，设置目标节点类型tgt_type、所有节点类型node_types和额外元路径extra_metapath。
    if args.dataset == 'DBLP':
        tgt_type = 'A'
        node_types = ['A', 'P', 'T', 'V']
        extra_metapath = []
    elif args.dataset == 'ACM':
        tgt_type = 'P'
        node_types = ['P', 'A', 'C']
        extra_metapath = []
    elif args.dataset == 'IMDB':
        tgt_type = 'M'
        node_types = ['M', 'A', 'D', 'K']
        extra_metapath = []
    elif args.dataset == 'Freebase':
        tgt_type = '0'
        node_types = [str(i) for i in range(8)]
        extra_metapath = []
    else:
        assert 0
    # 过滤extra_metapath，只保留长度大于args.num_hops + 1的路径。
    extra_metapath = [ele for ele in extra_metapath if len(ele) > args.num_hops + 1]

    # 打印当前特征传播的跳数args.num_hops。
    print(f'Current num hops = {args.num_hops} for feature propagation')

    # If out of memory when using gpu, please set `prop_device = 'cpu'`
    # 根据数据集和运行环境设置特征传播的设备prop_device和存储设备store_device。
    # prop_device: 如果数据集为Freebase且不使用CPU，则使用指定的GPU；否则使用CPU。
    # store_device: 设置存储设备为CPU。
    if args.dataset == 'Freebase':
        prop_device = f'cuda:{args.gpu}' if not args.cpu else 'cpu'
    else:
        prop_device = 'cpu'
    store_device = 'cpu'

    # 如果数据集为Freebase，则检查是否存在目录./Freebase_adjs，如果不存在则创建。并设置目标节点数量num_tgt_nodes。
    if args.dataset == 'Freebase':
        if not os.path.exists('./Freebase_adjs'):
            os.makedirs('./Freebase_adjs')
        num_tgt_nodes = dl.nodes['count'][0]

    # compute k-hop features
    raw_feats, data_size = neighboorAggregation(args, g, dl, extra_metapath, tgt_type, adjs, num_tgt_nodes, prop_device)

    # =======
    # 如果 args.amp 为真且不是在 CPU 上运行，则初始化 GradScaler，用于混合精度训练。
    # 否则，将 scalar 设为 None。
    if args.amp and not args.cpu:
        scalar = torch.cuda.amp.GradScaler()
    else:
        scalar = None

    # 克隆初始标签 init_labels。
    labels = init_labels.clone()

    # 选择设备：如果不是在 CPU 上运行，则使用指定的 GPU。
    device = f'cuda:{args.gpu}' if not args.cpu else 'cpu'
    # 对于非 IMDB 数据集，将标签转换为 long 类型，并移动到选定设备。
    # 对于 IMDB 数据集，将标签转换为 float 类型。
    if args.dataset != 'IMDB':
        labels_cuda = labels.long().to(device)
    else:
        labels = labels.float()
        labels_cuda = labels.to(device)

    # 循环遍历每一个种子值，设置随机种子以确保结果的可复现性。
    for seed in args.seeds:
        print('Restart with seed =', seed)
        args.seed = seed
        set_random_seed(args.seed)

        checkpt_folder = f'./output/{args.dataset}/'
        if not os.path.exists(checkpt_folder):
            os.makedirs(checkpt_folder)
        checkpt_file = checkpt_folder + uuid.uuid4().hex
        print('checkpt_file', checkpt_file)

        val_ratio = 0.2
        # 将训练验证节点（trainval_nid）复制并打乱，然后按比例划分为训练集和验证集。
        train_nid = trainval_nid.copy()
        np.random.shuffle(train_nid)
        split = int(train_nid.shape[0]*val_ratio)
        val_nid = train_nid[:split]
        train_nid = train_nid[split:]
        train_nid = np.sort(train_nid)
        val_nid = np.sort(val_nid)

        train_node_nums = len(train_nid)
        valid_node_nums = len(val_nid)
        test_node_nums = len(test_nid)
        trainval_point = train_node_nums
        valtest_point = trainval_point + valid_node_nums
        print(f'#Train {train_node_nums}, #Val {valid_node_nums}, #Test {test_node_nums}')

        labeled_nid = np.concatenate((train_nid, val_nid, test_nid))
        labeled_num_nodes = len(labeled_nid)
        num_nodes = dl.nodes['count'][0]

        # 检查是否有额外的未标注节点，如果存在，则记录这些节点。
        if labeled_num_nodes < num_nodes:
            flag = np.ones(num_nodes, dtype=bool)
            flag[train_nid] = 0
            flag[val_nid] = 0
            flag[test_nid] = 0
            extra_nid = np.where(flag)[0]
            print(f'Find {len(extra_nid)} extra nid for dataset {args.dataset}')
        else:
            extra_nid = np.array([])

        # clone raw_feats to avoid in-place modification for different seeds
        feats = {k: v.detach().clone() for k, v in raw_feats.items()}

        # =======
        # labels propagate alongside the metapath
        # =======
        # 根据 args.label_feats 参数，决定是否进行标签传播
        label_feats = {}
        if args.label_feats:
            if args.dataset != 'IMDB':
                label_onehot = torch.zeros((num_nodes, num_classes))
                label_onehot[train_nid] = F.one_hot(init_labels[train_nid], num_classes).float()
            else:
                label_onehot = torch.zeros((num_nodes, num_classes))
                label_onehot[train_nid] = init_labels[train_nid].float()

            if args.dataset == 'DBLP':
                extra_metapath = []
            elif args.dataset == 'IMDB':
                extra_metapath = []
            elif args.dataset == 'ACM':
                extra_metapath = []
            elif args.dataset == 'Freebase':
                extra_metapath = []
            else:
                assert 0

            # 过滤掉长度小于或等于 args.num_label_hops + 1 的元路径 extra_metapath。
            # 计算传播时使用的最大跳数 max_length。如果存在过滤后的元路径，max_length 将是 args.num_label_hops + 1 和 extra_metapath 中最长路径长度之间的较大值；否则，直接取 args.num_label_hops + 1。
            extra_metapath = [ele for ele in extra_metapath if len(ele) > args.num_label_hops + 1]
            if len(extra_metapath):
                max_length = max(args.num_label_hops + 1, max([len(ele) for ele in extra_metapath]))
            else:
                max_length = args.num_label_hops + 1

            # 输出当前用于标签传播的跳数（num_label_hops）。
            print(f'Current num hops = {args.num_label_hops} for label propagation')
            # compute k-hop feature
            # 这段代码的核心是处理在 Freebase 数据集上基于元路径的标签传播，包括加载或计算邻接矩阵的逻辑，特别是对 Freebase 数据集的优化处理。
            # 通过提前存储和加载矩阵，可以减少重复计算，提高效率。同时，代码还确保了加载的邻接矩阵与计算得到的矩阵一致性。
            # 记录开始标签传播的时间，用于后续计算传播所耗时间。
            prop_tic = datetime.datetime.now()
            # 当数据集为 Freebase 且 num_label_hops 小于等于 num_hops，并且没有额外的元路径时，直接从 raw_meta_adjs 中提取元路径邻接矩阵 meta_adjs。
            # 此处的 k[-1] == '0' 表示仅选择目标节点类型为 '0' 的邻接矩阵。
            if args.dataset == 'Freebase' and args.num_label_hops <= args.num_hops and len(extra_metapath) == 0:
                meta_adjs = {k: v for k, v in raw_meta_adjs.items() if k[-1] == '0' and len(k) < max_length}
            else:
                # 如果 Freebase 数据集且种子大于 0，并且已存在相应的保存文件，则从文件中加载元路径邻接矩阵 meta_adjs。
                if args.dataset == 'Freebase':
                    save_name = f'./Freebase_adjs/label_seed{args.seed}_hop{args.num_label_hops}'
                    if args.seed > 0 and os.path.exists(f'{save_name}_int64.npy'):
                        meta_adj_list = SparseAdjList(save_name, None, None, num_tgt_nodes, num_tgt_nodes, with_values=True)
                        meta_adjs = meta_adj_list.load_adjs(expand=True)
                    else:
                        # 否则，调用 hg_propagate_sparse_pyg 函数重新计算元路径邻接矩阵，并将其保存在指定位置。
                        meta_adjs = hg_propagate_sparse_pyg(
                            adjs, tgt_type, args.num_label_hops, max_length, extra_metapath, prop_feats=False, echo=True, prop_device=prop_device)
                        meta_adj_list = SparseAdjList(save_name, meta_adjs.keys(), meta_adjs, num_tgt_nodes, num_tgt_nodes, with_values=True)

                        # 验证重新加载的 meta_adjs 是否与内存中存储的邻接矩阵一致（通过 rowptr、col 和 value 进行比较）。
                        # 确认后，释放临时变量 tmp_adjs 和 tmp 并进行垃圾回收。
                        tmp = SparseAdjList(save_name, None, None, num_tgt_nodes, num_tgt_nodes, with_values=True)
                        tmp_adjs = tmp.load_adjs(expand=True)
                        for k in tmp.keys:
                            assert torch.all(meta_adjs[k].storage.rowptr() == tmp_adjs[k].storage.rowptr())
                            assert torch.all(meta_adjs[k].storage.col() == tmp_adjs[k].storage.col())
                            assert torch.all(meta_adjs[k].storage.value() == tmp_adjs[k].storage.value())
                        del tmp_adjs, tmp
                        gc.collect()
                else:
                    # 非 Freebase 的情况下，直接调用 hg_propagate_sparse_pyg 函数计算 meta_adjs。
                    meta_adjs = hg_propagate_sparse_pyg(
                        adjs, tgt_type, args.num_label_hops, max_length, extra_metapath, prop_feats=False, echo=True, prop_device=prop_device)

            # 打印每个元路径邻接矩阵的键名和尺寸信息。
            print(f'For label propagation, meta_adjs: (in SparseTensor mode)')
            for k, v in meta_adjs.items():
                print(k, v.sizes())
            print()

            # 这段代码的主要功能是执行标签传播，并根据传播结果更新 label_feats 字典。代码还包括了对不同数据集的特殊处理，特别是 Freebase 数据集。
            if args.dataset == 'Freebase':
                if False:
                    label_onehot_g = label_onehot.to(prop_device)
                    for k, v in tqdm(meta_adjs.items()):
                        if args.dataset != 'Freebase':
                            label_feats[k] = remove_diag(v) @ label_onehot
                        else:
                            label_feats[k] = (remove_diag(v).to(prop_device) @ label_onehot_g).to(store_device)

                    del label_onehot_g
                    if not args.cpu: torch.cuda.empty_cache()
                    gc.collect()

                    condition = lambda ra,rb,rc,k: rb > 0.2
                    check_acc(label_feats, condition, init_labels, train_nid, val_nid, test_nid, show_test=False)

                    left_keys = ['00', '000', '0000', '0010', '0030', '0040', '0050', '0060', '0070']
                    remove_keys = list(set(list(label_feats.keys())) - set(left_keys))
                    for k in remove_keys:
                        label_feats.pop(k)
                else:
                    # 定义了一组保留的键 left_keys，这些键代表不同的元路径。
                    # 移除 meta_adjs 中不在 left_keys 中的键，减少后续计算的复杂度。
                    left_keys = ['00', '000', '0000', '0010', '0030', '0040', '0050', '0060', '0070']
                    remove_keys = list(set(list(meta_adjs.keys())) - set(left_keys))
                    for k in remove_keys:
                        meta_adjs.pop(k)

                    # 将 label_onehot 标签矩阵移至 prop_device，以便在后续的矩阵操作中提高效率。
                    label_onehot_g = label_onehot.to(prop_device)
                    # 对于每个 meta_adjs 中的邻接矩阵 v，执行标签传播计算。
                    for k, v in tqdm(meta_adjs.items()):
                        if args.dataset != 'Freebase':
                            # remove_diag(v) 用于移除矩阵的对角线元素（即自连接），然后与标签矩阵 label_onehot_g 相乘，得到新的标签特征 label_feats[k]。
                            label_feats[k] = remove_diag(v) @ label_onehot
                        else:
                            # 如果数据集是 Freebase，传播操作将在 prop_device 上执行，并且结果移动到 store_device。
                            label_feats[k] = (remove_diag(v).to(prop_device) @ label_onehot_g).to(store_device)

                    # 删除 label_onehot_g，释放 GPU 内存。
                    # 如果未使用 CPU，调用 torch.cuda.empty_cache() 清理 GPU 内存。
                    # 使用 gc.collect() 进行垃圾回收，清理无用的内存。
                    del label_onehot_g
                    if not args.cpu: torch.cuda.empty_cache()
                    gc.collect()
            else:
                # 其他数据集的标签传播处理
                # 这部分代码与 Freebase 处理类似，只是在不同数据集时采取不同的设备迁移策略。
                for k, v in tqdm(meta_adjs.items()):
                    # todo 这个地方可以优化
                    if args.dataset != 'Freebase':
                        label_feats[k] = remove_diag(v) @ label_onehot
                    else:
                        label_feats[k] = (remove_diag(v).to(prop_device) @ label_onehot_g).to(store_device)
                gc.collect()

                # 对于不同的数据集，使用 check_acc 函数评估标签传播的结果。
                # IMDB 数据集使用二元交叉熵损失（bce），且不显示测试集结果（show_test=False）；其他数据集则显示测试集结果。
                if args.dataset == 'IMDB':
                    condition = lambda ra,rb,rc,k: True
                    check_acc(label_feats, condition, init_labels, train_nid, val_nid, test_nid, show_test=False, loss_type='bce')
                else:
                    condition = lambda ra,rb,rc,k: True
                    check_acc(label_feats, condition, init_labels, train_nid, val_nid, test_nid, show_test=True)
            # 输出参与标签传播的元路径键名，供用户参考。
            print('Involved label keys', label_feats.keys())

            # label_feats = {k: v[init2sort] for k,v in label_feats.items()}
            # 记录并输出标签传播的时间，便于了解计算效率。
            prop_toc = datetime.datetime.now()
            print(f'Time used for label prop {prop_toc - prop_tic}')

        # =======
        # Train & eval loaders
        # 这段代码的目的是准备用于训练和评估的批数据加载器 (train_loader, eval_loader, full_loader) 。
        # 它还涉及了一些与掩码（mask）和标签平滑（smoothing）相关的处理，虽然在当前代码中未完全实现或被注释掉了。
        # =======
        # 使用 torch.utils.data.DataLoader 为训练数据创建一个加载器。
        # train_nid 包含训练节点的ID。
        # batch_size 是训练时使用的批大小。
        # shuffle=True 表示每个epoch都会随机打乱数据顺序。
        # drop_last=False 表示不会丢弃最后一个不足一个批的数据。
        train_loader = torch.utils.data.DataLoader(
            train_nid, batch_size=args.batch_size, shuffle=True, drop_last=False)

        # =======
        # Mask & Smooth
        # with_mask 被设置为 False，表示当前不会应用掩码。
        # 注释的代码中，with_mask 本来是用来处理 Freebase 数据集的掩码，其他数据集未实现。
        # 如果 with_mask=True，会创建 train_mask 和 full_mask，用于训练和完整数据的掩码。
        # =======
        with_mask = False
        # if args.dataset == 'Freebase':
        #     init_mask = {k: v.storage.rowcount() != 0 for k, v in feats.items()}
        #     with_mask = True
        # else:
        #     print(f'TODO: `with_mask` has not be implemented for {args.dataset}')

        # if with_mask:
        #     train_mask = {k: (v[:labeled_num_nodes] & (torch.randn(labeled_num_nodes) > 0)).float() for k, v in init_mask.items()}
        #     full_mask = {k: v.float() for k, v in init_mask.items()}
        # else:
        #     train_mask = full_mask = None

        # with_smooth = False

        # Freebase train/val/test/full_nodes: 1909/477/5568/40402
        # IMDB     train/val/test/full_nodes: 1097/274/3202/359
        # 初始化了两个空列表 eval_loader 和 full_loader，用于存储评估和完整数据的批次。
        # 评估的批大小 (eval_batch_size) 设置为训练批大小的两倍。
        eval_loader, full_loader = [], []
        eval_batch_size = 2 * args.batch_size

        # 遍历 labeled_nid，将它们分成批次，添加到 eval_loader 中。
        # 每个批次都包含节点的特征 (batch_feats)、标签特征 (batch_labels_feats) 和可能的掩码 (batch_mask)。
        # 这里的掩码只在 with_mask=True 时才会创建。
        for batch_idx in range((labeled_num_nodes-1) // eval_batch_size + 1):
            batch_start = batch_idx * eval_batch_size
            batch_end = min(labeled_num_nodes, (batch_idx+1) * eval_batch_size)
            batch = torch.LongTensor(labeled_nid[batch_start:batch_end])

            batch_feats = {k: x[batch] for k, x in feats.items()}
            batch_labels_feats = {k: x[batch] for k, x in label_feats.items()}
            if with_mask:
                batch_mask = {k: x[batch] for k, x in full_mask.items()}
            else:
                batch_mask = None
            eval_loader.append((batch, batch_feats, batch_labels_feats, batch_mask))

        # 类似于 eval_loader，full_loader 则是用于处理 extra_nid 中额外节点的批次。
        for batch_idx in range((len(extra_nid)-1) // eval_batch_size + 1):
            batch_start = batch_idx * eval_batch_size
            batch_end = min(len(extra_nid), (batch_idx+1) * eval_batch_size)
            batch = torch.LongTensor(extra_nid[batch_start:batch_end])

            batch_feats = {k: x[batch] for k, x in feats.items()}
            batch_labels_feats = {k: x[batch] for k, x in label_feats.items()}
            if with_mask:
                batch_mask = {k: x[batch] for k, x in full_mask.items()}
            else:
                batch_mask = None
            full_loader.append((batch, batch_feats, batch_labels_feats, batch_mask))

        # =======
        # Construct network
        # =======
        # 首先清理了 CUDA 缓存和 Python 的垃圾回收。
        if not args.cpu: torch.cuda.empty_cache()
        gc.collect()

        # 构建了 SeHGNN 模型，并将其移动到指定设备（CPU 或 GPU）。
        model = SeHGNN(args.dataset, args.embed_size, args.hidden, num_classes, feats.keys(), label_feats.keys(), tgt_type,
                       args.dropout, args.input_drop, args.att_drop, args.n_fp_layers, args.n_task_layers, args.act,
                       args.residual, data_size=data_size)
        model = model.to(device)
        # 打印模型结构及其参数总数（在 args.seed 为初始值时）。
        if args.seed == args.seeds[0]:
            print(model)
            print('# Params:', get_n_params(model))

        # 对于 IMDB 数据集，使用二元交叉熵损失 BCEWithLogitsLoss。对于其他数据集，使用交叉熵损失 CrossEntropyLoss。
        if args.dataset == 'IMDB':
            loss_fcn = nn.BCEWithLogitsLoss()
        else:
            loss_fcn = nn.CrossEntropyLoss()
        # 使用 Adam 优化器，学习率和权重衰减从参数中读取。
        optimizer = torch.optim.Adam(model.parameters(), lr=args.lr,
                                    weight_decay=args.weight_decay)

        # 初始化了多个变量以追踪最佳模型状态、训练时间以及各个 epoch 的损失和准确性。
        best_epoch = -1
        best_val_loss = 1000000
        best_test_loss = 0
        best_val = (0,0)
        best_test = (0,0)
        val_loss_list, test_loss_list = [], []
        val_acc_list, test_acc_list = [], []
        actual_loss_list, actual_acc_list = [], []
        store_list = []
        best_pred = None
        count = 0

        train_times = []

        for epoch in tqdm(range(args.epoch)):
            # 每个 epoch 的训练时间、损失和准确度都会被记录。
            gc.collect()
            if not args.cpu: torch.cuda.synchronize()
            start = time.time()
            # 用 train 函数进行模型训练，之后计算并存储训练时间。
            loss, acc = train(model, feats, label_feats, labels_cuda, loss_fcn, optimizer, train_loader, evaluator, scalar=scalar)
            if not args.cpu: torch.cuda.synchronize()
            end = time.time()

            log = f'Epoch {epoch}, training Time(s): {end-start:.4f}, estimated train loss {loss:.4f}, acc {acc[0]*100:.4f}, {acc[1]*100:.4f}\n'
            if not args.cpu: torch.cuda.empty_cache()
            train_times.append(end-start)

            # 在验证集和测试集上评估模型性能。
            # 记录验证损失和测试损失，并基于它们更新最佳模型。
            start = time.time()
            with torch.no_grad():
                model.eval()
                raw_preds = []
                for batch, batch_feats, batch_labels_feats, batch_mask in eval_loader:
                    batch = batch.to(device)
                    batch_feats = {k: x.to(device) for k, x in batch_feats.items()}
                    batch_labels_feats = {k: x.to(device) for k, x in batch_labels_feats.items()}
                    if with_mask:
                        batch_mask = {k: x.to(device) for k, x in batch_mask.items()}
                    else:
                        batch_mask = None
                    raw_preds.append(model(batch, batch_feats, batch_labels_feats, batch_mask).cpu())

                raw_preds = torch.cat(raw_preds, dim=0)
                loss_train = loss_fcn(raw_preds[:trainval_point], labels[train_nid]).item()
                loss_val = loss_fcn(raw_preds[trainval_point:valtest_point], labels[val_nid]).item()
                loss_test = loss_fcn(raw_preds[valtest_point:labeled_num_nodes], labels[test_nid]).item()

            if args.dataset != 'IMDB':
                preds = raw_preds.argmax(dim=-1)
            else:
                preds = (raw_preds > 0.).int()

            train_acc = evaluator(preds[:trainval_point], labels[train_nid])
            val_acc = evaluator(preds[trainval_point:valtest_point], labels[val_nid])
            test_acc = evaluator(preds[valtest_point:labeled_num_nodes], labels[test_nid])

            end = time.time()
            log += f'evaluation Time: {end-start:.4f}, Train loss: {loss_train:.4f}, Val loss: {loss_val:.4f}, Test loss: {loss_test:.4f}\n'
            log += f'Train acc: ({train_acc[0]*100:.4f}, {train_acc[1]*100:.4f}), Val acc: ({val_acc[0]*100:.4f}, {val_acc[1]*100:.4f})'
            log += f', Test acc: ({test_acc[0]*100:.4f}, {test_acc[1]*100:.4f})\n'

            # 如果验证损失（或准确率）比之前更好，则更新最佳模型并保存其状态。
            if (args.dataset != 'Freebase' and loss_val < best_val_loss) or (args.dataset == 'Freebase' and sum(val_acc) > sum(best_val)):
                best_epoch = epoch
                best_val_loss = loss_val
                best_test_loss = loss_test
                best_val = val_acc
                best_test = test_acc

                best_pred = raw_preds
                torch.save(model.state_dict(), f'{checkpt_file}.pkl')

            if epoch - best_epoch > args.patience:
                break

            if epoch > 0 and epoch % 10 == 0: 
                log = log + f'\tCurrent best at epoch {best_epoch} with Val loss {best_val_loss:.4f} ({best_val[0]*100:.4f}, {best_val[1]*100:.4f})' \
                    + f', Test loss {best_test_loss:.4f} ({best_test[0]*100:.4f}, {best_test[1]*100:.4f})'
            print(log)

        print('average train times', sum(train_times) / len(train_times))

        print(f'Best Epoch {best_epoch} at {checkpt_file.split("/")[-1]}\n\tFinal Val loss {best_val_loss:.4f} ({best_val[0]*100:.4f}, {best_val[1]*100:.4f})'
            + f', Test loss {best_test_loss:.4f} ({best_test[0]*100:.4f}, {best_test[1]*100:.4f})')

        # 在整个数据集上运行模型，生成最终的预测并保存。
        all_pred = torch.empty((num_nodes, num_classes))
        all_pred[labeled_nid] = best_pred
        if len(full_loader):
            model.load_state_dict(torch.load(f'{checkpt_file}.pkl', map_location='cpu'), strict=True)
            if not args.cpu:
                torch.cuda.empty_cache()
            with torch.no_grad():
                model.eval()
                raw_preds = []
                for batch, batch_feats, batch_labels_feats, batch_mask in full_loader:
                    batch = batch.to(device)
                    batch_feats = {k: x.to(device) for k, x in batch_feats.items()}
                    batch_labels_feats = {k: x.to(device) for k, x in batch_labels_feats.items()}
                    if with_mask:
                        batch_mask = {k: x.to(device) for k, x in batch_mask.items()}
                    else:
                        batch_mask = None
                    raw_preds.append(model(batch, batch_feats, batch_labels_feats, batch_mask).cpu())
                raw_preds = torch.cat(raw_preds, dim=0)

            all_pred[extra_nid] = raw_preds
        torch.save(all_pred, f'{checkpt_file}.pt')

        # 根据模型输出的预测概率生成最终的预测标签，并将这些标签保存到指定文件中。
        if args.dataset != 'IMDB':
            predict_prob = all_pred.softmax(dim=1)
        else:
            predict_prob = torch.sigmoid(all_pred)

        test_logits = predict_prob[test_nid]
        if args.dataset != 'IMDB':
            pred = test_logits.cpu().numpy().argmax(axis=1)
            dl.gen_file_for_evaluate(test_idx=test_nid, label=pred, file_name=f"{args.dataset}_{args.seed}_{checkpt_file.split('/')[-1]}.txt")
        else:
            pred = (test_logits.cpu().numpy()>0.5).astype(int)
            dl.gen_file_for_evaluate(test_idx=test_nid, label=pred, file_name=f"{args.dataset}_{args.seed}_{checkpt_file.split('/')[-1]}.txt", mode='multi')

        # 计算并打印训练、验证和测试集的准确率，以便评估模型在各个数据集上的表现。
        if args.dataset != 'IMDB':
            preds = predict_prob.argmax(dim=1, keepdim=True)
        else:
            preds = (predict_prob > 0.5).int()
        train_acc = evaluator(preds[train_nid], labels[train_nid])
        val_acc = evaluator(preds[val_nid], labels[val_nid])
        test_acc = evaluator(preds[test_nid], labels[test_nid])

        print(f'train_acc ({train_acc[0]*100:.2f}, {train_acc[1]*100:.2f}) ' \
            + f'val_acc ({val_acc[0]*100:.2f}, {val_acc[1]*100:.2f}) ' \
            + f'test_acc ({test_acc[0]*100:.2f}, {test_acc[1]*100:.2f})')
        print(checkpt_file.split('/')[-1])

        del model
        if not args.cpu: torch.cuda.empty_cache()


if __name__ == '__main__':
    args = parse_args()

    if args.dataset == 'ACM':
        # We do not keep the node type `field`(F) as raw features represent exactly the distribution of nearby field nodes
        args.ACM_keep_F = False

    print(args)
    main(args)
